refactor(kv-cache): one declaration per fact about the KV pools - #2147

Every consumer takes max(floor, rows), so the floor only adds split-K where
the grid is too small to fill the machine and is an identity everywhere
else. The splits are numerically inert: each CTA owns a disjoint KV-column
range, with no cross-CTA partial sums.

512 left the machine idle on long contexts, which is exactly where rows are
scarce -- the logits buffer widens as the context grows, so the Q chunk and
with it the row count shrink. Decode rows=128 W~32768 54.1us -> 51.2us;
prefill rows=1024 224.6us -> 206.2us.

4096 is not any single shape's optimum, because CTA-count quantization makes
the ordering shape-specific, but it has the smallest worst-case regret of
the values tried. Worth re-tuning against a real workload mix.

The comment at the prefill call site said the floor was for decode. It is
not: a long-context prefill is short on rows for the same reason.

… serves

`StateArena` was never state-specific. Its `buf=` / `slot_stride=` /
`live_entries=` hooks already let DeepSeek-V4 carve it out of a plane it
shares with paged blocks, so its only tie to the STATE pool was the name --
and a paged KV pool is about to read the same field declarations through a
sibling materializer. What separates the two is where the entry axis sits,
which its own docstring already said ("laid out entry-major"), so that is
what it is called now. Left as it was, the pair would have read as "the
state one" against "the by-layer one", implying the first cannot serve PAGE.

`entry` is not a new word for this: `sub_pool_spec` already sizes both pools
in entries -- `entry_bytes`, `entries_per_req`, "entry class" -- because it
is the unit-agnostic term. A PAGE entry is a block, a STATE entry is one
request's state.

  StateField      -> EntryField
  StateArena      -> EntryMajorArena
  SplitStateArena -> SplitEntryMajorArena
  state_arena.py  -> entry_arena.py

Splitting the declarations out of the arena was tried and reverted, and the
module docstring now records why, because it is the obvious thing to try
next: the `pool_layout` guard bans relative imports too (`imported_roots`
scores a relative import as reaching `atom`), so every member of the package
is a leaf, and a field list cannot live apart from an arena over it.

No behaviour change. Bodies are byte-identical modulo the renames except for
three strings that name the old type.

Seven places spelled out the SHUFFLE shapes of a paged MHA block. They agreed
on every byte and disagreed on a view: the Eagle3 draft's builder declared V
non-transposed where the others declared it 5-D, and `_bind_kv_cache_to_modules`
did the same. Nothing could catch it -- the byte count is identical either way,
and the only consumer that cares reads `v_cache.ndim` at runtime, densifying
anything that is not already 5-D. On MiniMax-M3 that relaid the whole 1.4 GiB
draft pool once per chunked-prefill forward: 7.1% of GPU time, 61% of what the
draft model cost.

`MhaKvPool` is that declaration and the two arenas over it. It takes five
numbers -- layers, block_size, num_kv_heads, head_dim, kv_dtype -- and never
sees a ModelRunner, a config or a forward_vars, which is what makes a second
one askable for. That was the actual obstacle: a backend also owns per-step
metadata, and `CommonAttentionBuilder.__init__` ends by writing into
`model_runner.forward_vars`, so building a second builder to get a second pool
would have overwritten the target's buffers. Splitting the two jobs is what
lets the draft have a pool without anyone copying shapes.

So a draft no longer carries an attention. `Eagle3DraftBuilder` hardcoded MHA;
`DraftKvBuilder` is handed a pool and keeps only the part that is about being a
draft -- its blocks are the target's blocks, re-paged at its own block size.
Which pool it gets comes from the draft's own config through the same selector
the target uses, and the backend answers `make_kv_pool` for its own flavor
(None means "a draft of this flavor binds into the target's pool", which is
MLA and the default). Adding a flavor is one method on that backend; nothing in
spec_decode imports an attention any more.

Sizing derives from the same declaration, so a block cannot be charged for at
one size and built at another. M3's `entry_bytes` is 3502080 before and after.

Also folded in:
  - `LayerMajorArena`, the sibling of `EntryMajorArena` for a pool whose layer
    axis is outermost. It has no `entry(i)` on purpose: a block's bytes are
    `blocks` apart, so nothing can copy or register one as a unit. That is the
    deletion condition too -- turn the pool block-major and one class serves
    both.
  - `TritonMHAMetadataBuilder.build_kv_cache_tensor` shrinks to the parent's
    bind plus the layout flag, which was the only thing it added.
  - `GDNAttentionMetadataBuilder` allocated its own tensor but inherited the
    bind, so it needed a pool as well; its shorter layer axis is now one
    `_kv_pool_layers` hook that sizing and allocation share.
  - `_bind_kv_cache_to_modules` runs the builders' own bind loop instead of a
    second copy of it, after `adopt_imported_kv_pool` re-derives the pool over
    the IPC-imported buffers. Handing those buffers to an arena is also the
    check that they fit the declaration.

Verified on MiniMax-M3-MXFP4 + EAGLE3, tp8 unless noted:
  - trace (tp4, ISL 100k): float8_copy 416 -> 0 on all four ranks, with
    cp_mha_gather_cache still at 104, so the path runs rather than being
    skipped.
  - GSM8K 5-shot chat: 0.9431 flexible / 0.9439 strict against a 0.9469
    baseline and a 0.93 threshold. Run with --level 0, which the CI entry does
    not use, so it is a band check rather than a point comparison.
  - MTP acceptance 73.55% over 123k draft tokens against a 0.715 threshold --
    the sensitive one here, since a draft KV write the verifier reads back
    differently collapses it rather than nudging it.
  - Unit tests unchanged at 4 pre-existing failures, all aiter-side.

Not exercised: P/D, MiMo-V2, MLA targets. GDN was caught by reading the
inheritance chain, not by a test.

`MhaKvPool` left MLA as the asymmetric half, and MLA's copy of the problem was
worse: `sub_pool_specs`, `allocate_kv_cache_tensors` and `build_kv_cache_tensor`
each wrote the literal `576`, while `mla_kv_entry_dim` -- the expression that
exists precisely because a NoPE model declares a padded width that need not
equal `kv_lora_rank + qk_rope_head_dim` -- was read only by `mla_qk_head_dim`
and the Kimi backend. Every model in the tree happens to be 576, so it is a
latent split rather than a live one; it is closed here as a substitution that
is provably an identity today, not as a bugfix.

`MlaKvPool` is that declaration. MLA packs k_c and k_pe into one row and
absorbs V, so a block is one field rather than four, plus an indexer field the
sparse variants add -- owned by some layers only, and holding one row per token
or, where the indexer pools, fewer. Differing layer counts and row counts are
what `EntryField` already says, so both are fields and neither needs a second
pool.

An entry is the *scheduler* block, which is what `page_pool` charges. MLA
usually pages at size 1, so `block_ratio` rows of the allocation make one
entry; keeping the entry at the scheduler block is what lets a row of
`region_tensors` be the unit a transfer registers. Three things fall out:

  - `get_kv_transfer_tensors` drops its `* block_ratio`, because `stride(0)`
    is already the bytes a block owns;
  - `_index_cache_block_bytes` goes, base and Kimi override both, since the
    field's own arithmetic answers it;
  - `runner.physical_block_size` goes. It was the builder's own `block_size`
    mirrored onto the runner, with five readers, all of them inside builders.

`kimi_mla_gdn` gets the same pool over a shorter layer axis (`_kv_pool_layers`),
and `PageUnitGeometryMixin` reads the pool's view through a new
`_page_unit_kv_cache` hook, symmetric with the index hook it already had.

The 39 tests this broke were all real: MLA's fixtures pinned the allocator's
call shape and the tensor doubles it sliced. They now assert the views a reader
binds, and the doubles shrank to a fake pool -- `layer(name, i)` and
`region_tensors()` are a much narrower interface than a tensor that has to
behave like a tensor. One of them earned its keep immediately, catching that
the transfer path must use the layer count the pool was *built* with: a hybrid
caches for fewer layers than the model has.

Also here, found by the M3 gate rather than by tests: the block count a pool is
built at is `config.num_kvcache_blocks`, not `pool_plan.paged_entries`. Sizing
runs in every runner subprocess and they disagree by a few blocks, so EngineCore
takes one answer and broadcasts it; the plan is this rank's own estimate. Using
it made the KV pool larger than the sparse index cache sized alongside it, and
`fused_qknorm_idxrqknorm` refused the mismatch at capture time. The comment says
so at the one place that reads it.

`wait_server_ready.sh` learns `openai_server.py: error:` -- argparse refusing
the command line kills the process in under a second, and the poll used to wait
out its whole budget. A bare `Traceback` was tried and reverted: torch logs
whole tracebacks at WARNING during Inductor warmup, hundreds of them, and it
failed a healthy start on first use.

Verified:
  - MLA byte parity, measured both ways: Kimi-K3 `entry_bytes` is 1769472 on
    this commit and on HEAD, and R1-0528's 571392 decomposes as 62 layers x 16
    tokens x 576 x 1B, which is the old formula's expression.
  - DeepSeek-R1-0528-MXFP4 MTP3 tp4: GSM8K 0.9507 flexible against a 0.93
    threshold, MTP acceptance 64.74% against 0.593 over 150k draft tokens.
  - MiniMax-M3 + EAGLE3 re-run after the `physical_block_size` removal:
    float8_copy still 0 on all four ranks, cp_mha_gather_cache still 104,
    `entry_bytes` still 3502080, throughput flat.
  - Unit tests unchanged at 4 pre-existing failures, all aiter-side.

Not exercised: Kimi-K3 accuracy. It scores 0 with `!!!!` output on this box,
and reproduces identically on HEAD down to the first token, so it is upstream
of this change -- the aiter working tree is mid-rebuild. P/D is untested as
before.

…s block count

Two changes, in one commit because they interleave inside
`AiterAttentionMetadataBuilder.allocate_kv_cache_tensors`; splitting by file
would leave a first commit that does not import.

1. MiniMax-M3's indexer key cache becomes a field of `MhaKvPool`

It was the last hand-written declaration on the MHA side, and it was written
twice: a formula added to `block_bytes` in `sub_pool_specs`, and a
`torch.zeros` in `allocate_kv_cache_tensors`. The two even read different
block-count attributes -- exactly the shape that produced

  fused_qknorm_idxrqknorm.cu:851 index_cache must contain at least
  num_blocks * block_size * 128 elements

when one of them moved. Now `index_layers` / `index_dim` / `index_dtype` are
constructor arguments, symmetric with `MlaKvPool.index`, so one expression
answers sizing and allocation and a divergence is unrepresentable.

Three things fall out, none of them cosmetic:

- The old sizing charged `sparse_layers * backend_page * dim` while the
  allocation was `blocks * block_ratio` pages deep, i.e. it undercharged by
  `block_ratio`. A field is declared per scheduler block, so the count is now
  right for any ratio. Every shipped M3 config runs the two block sizes equal,
  which is why no budget ever came up short.
- `runner.sparse_attention_index_cache` is gone rather than re-created. Nothing
  outside this builder read it, and neither `ModelRunner.release` nor the
  rollout memory manager named it -- so it was a bare attribute that survived a
  rollout sleep. The pool is now its only owner and `MhaKvPool.release` frees
  it.
- `_sparse_attention_topk_cache_state` moves out from under the indexer gate.
  It is a per-run memo keyed off `use_index_cache`; being nested under the
  cache claimed a dependency that does not exist.

P/D now refuses this model instead of guessing: the IPC export ships `kv_cache`
and `kv_scale` only, so a decode side has no imported buffer for the indexer
cache, and allocating a fresh one would give it zeroed keys and a top-k that
picks the wrong blocks without saying so.

2. The block count arrives as an argument

`ModelRunner` used to reach into the builder for `block_ratio`, multiply, hang
the product back on itself as `num_physical_kvcache_blocks`, and let five
backends read it off again -- two of them reading `config.num_kvcache_blocks`
instead, for the same quantity. Sizing runs per subprocess and the answers
differ (48425 / 48339 / 48310 on M3 tp4) while EngineCore broadcasts one, so
reading the wrong attribute builds one pool at one count and something sized
off another at a second.

`allocate_kv_cache_tensors` now takes a keyword-only `blocks`, the same
argument `adopt_imported_kv_pool` already took, and each builder records
`self.num_blocks = blocks * self.block_ratio` -- its own page, which is what
its kernels index. `num_physical_kvcache_blocks` is gone.

The units are the point, and they are kept apart even where they are the same
number: `blocks`, a pool entry and `req.block_ids` are the scheduler's;
`num_blocks` is the backend's. Conversions are written out; none is elided
because a ratio happens to be 1.

`KVTransferTensors.num_blocks` becomes `field(init=False)`. Three backends used
to state it themselves, in two different units, while the field's consumers
address `req.block_ids`. It is now set once by `ModelRunner.set_kv_cache_data`'s
caller -- after the draft appends its regions, and the only place holding both
the complete region list and the scheduler's count -- and `set_block_count`
checks every region divides into exactly that many units. The failure it
catches has identical total bytes: a region registered per token rather than
per scheduler block has `block_ratio` times the units at `1/block_ratio` the
stride, so both ends agree on the address arithmetic and disagree on the
stride. No byte count distinguishes them.

Tests

`tests/test_kv_pool_contract.py` is new: an AST check that every implementation
of the allocate hook takes `blocks` keyword-only with no default. A source grep
would have asserted on the text of the code it tests; the signature is the
invariant that makes the runner attribute unnecessary. Also 8 cases pinning the
indexer field, 4 pinning the block id space, and one pinning that an indexer
region is one scheduler block wide -- the single place the fold is deliberately
not byte-for-byte with what it replaced.

`tests/test_mla_index_cache.py`'s allocate fixture now runs at `block_ratio`
16, not 1: at 1 it could not tell the two counts apart.

Verification

- unit tests 5443 passed / 4 failed, against a 5415 / 4 baseline; the 4 are
  pre-existing aiter-side failures (`_top_k_per_row_decode` signature,
  `_is_gfx1250`), neither of which CI runs
- `entry_bytes` byte parity, offline and in the server log: 2961408 == 2961408
  for the target pool, `sub-pool kv: entry_bytes=3502080` unchanged
- MiniMax-M3-MXFP4 + EAGLE3, tp4: GSM8K 0.9363 and 0.9431 over two runs,
  acceptance 75.37% / 73.00%. Five runs on this branch span 0.9363-0.9462 and
  0.9431 appears in three of them, across both this change and the one before
  it, so the arms overlap
- DeepSeek-V4-Flash-DSpark, tp1: GSM8K 0.9484, acceptance 65.29%, 4.26
  toks/fwd -- the recorded baseline on all three
- revert checks: swapping arena order reddens the region-order test; removing
  GDN's `blocks` parameter reddens that implementation's two contract cases and
  nothing else

Not exercised: P/D (`set_block_count` has unit coverage only), MiMo-V2, and
Kimi-K3, whose accuracy is blocked on an unrelated aiter-side fault.
`tests/test_mla_index_cache.py` skips at module level without aiter, so CI
collects it and runs none of it.

The target's KV, its dequantization scales, MiniMax-M3's indexer keys and an
EAGLE3 draft's sibling pool were up to five `torch.zeros` under five runner
attributes, while the budget had long treated them as one entry class -- two
specs naming the same class sum, which is how a draft rides the target's block
ids. Sizing said one pool; memory was five.

They are one buffer now. `ModelRunner._carve_paged_pool` asks each builder how
many bytes its pool takes, places them with `plan_regions`, and hands each its
region. `paged_pool_bytes` answering zero means "I allocate my own", which is
where DeepSeek-V4 stays: half of what its PAGE entry is priced for lives in
the plane pool `allocate_per_req_cache` makes later.

What that buys:

  - `allocate_kv_cache_tensors(..., buf=)` and `adopt_imported_kv_pool(blocks,
    buf)` collapse into one expression, `pool.allocate(blocks, dev, buf=...)`.
    The P/D decode side carves with the same walk instead of being told the
    offsets, so the two ends cannot disagree about where a region starts.
  - The IPC export ships one handle instead of two, and the sparse indexer
    cache travels inside it. The refusal in the MHA adopt path went with that:
    its reason -- "the export ships kv_cache and kv_scale and nothing else" --
    stopped being true.
  - `runner.kv_scale` and `runner.index_cache` are gone, with their entries in
    both release lists. `kv_cache` is the pool.
  - MiMo-V2's per-module pools take successive slices of their builder's
    region, so they are exported and freed like every other model's.

Each pool carries its groups as one list that the price and the allocation
both walk (`field_groups` -> `entry_bytes` / `carve_layer_major`), so a fourth
group cannot be added to one and missed by the other.

The one thing checked on every start: a builder either allocates its own pool
or takes exactly what its own PAGE spec was charged. A region sized off a
second reading of the layout surfaces as a wrong answer while serving, never
as a bad number at startup.

Unit tests 5460 passed / 4 failed, the four being aiter-side failures already
on HEAD. MiniMax-M3 tp4: `entry_bytes=3502080`, byte-identical to before;
GSM8K 0.9454 and 0.9447 at 73.28% / 73.08% MTP acceptance, against five
earlier runs on this branch spanning 0.9363-0.9462; trace
`float8_copy_kernel_cuda` 0 on all four ranks and `cp_mha_gather_cache_kernel`
104. DeepSeek-V4-Flash-DSpark tp2, which opts out: 0.9507, 65.12%, 4.26
toks/fwd against a 0.9484 / 65.09% / 4.25 baseline.

Not exercised: P/D (the single-handle format has unit coverage only), MiMo-V2,
Kimi-K3.

…ules declare the pools

A layer's row in a pool was recovered from `layer_id` by knowing the model's
shape: `layer_id // full_attention_interval` for a Qwen3-Next full-attention
layer, its complement for a linear one, `num_full_attn + (layer_id -
mtp_start)` past a draft's boundary, a dict keyed by layer for Kimi's two
kinds. Four spellings of a counter, each carrying an assumption about a model
it did not name -- and `layer_id` is already a counter, of every layer the
walk registered.

It is `take_slot(kind)` now: a module's row is its position among the modules
of its kind, which is the order the walk visits them in. The runner clears the
counters before each walk, because the walk is the runner's -- a pool is bound
again on a P/D import and on a rollout wake, and that also fixes a draft
builder whose own counter was never reset (a wake would have handed its
one-layer pool row 1).

The pools themselves come from the modules too. A module says what it caches,
so `MhaKvPool`s are the modules counted by `(num_kv_heads, head_dim)`: one for
almost every model, two for MiMo-V2-Flash, whose sliding-window layers carry a
second head count -- a second geometry in one block-id space, which is what
sizing had been summing for it all along. Which layers of a hybrid are full
attention, where a draft's stack begins, what a PP stage holds: all of it is
now a property of what modules exist, and one predicate answers it for the
sizing walk and the bind walk both, so a layer cannot be charged for and left
unbound.

`is_mimo_v2` is gone from the tree with the last of the per-layer allocation
it gated -- and with it MiMo-V2's cursor into its own region, its
`_kv_layer_cache_store`, the P/D refusal, the transfer path's second branch,
and the Triton backend's refusal, which named that allocation path. The
GDN-hybrid builder loses both of its pool overrides: counting the modules it
owns leaves the linear-attention layers out without it saying so. Also gone,
having lost their last reader: `runner.mtp_start_layer_idx`,
`runner.full_attention_interval`, `runner.kda_attention_layers`, Kimi's two
layer-index dicts, and `_kv_pool_layers` on the MHA side.

`entry_arena.carve` is the placement rule the runner's pool, a pool's field
groups and a builder's several pools now share.

Unit tests 5460 passed / 4 failed, the four being aiter-side failures already
on HEAD. `tests/test_bind_walk_slots.py` pins the counter against the deleted
formulas on the shape they were written for.

MiniMax-M3 tp4: `entry_bytes=3502080` unchanged, GSM8K 0.9401 at 73.14%
acceptance. DeepSeek-V4-Flash-DSpark tp2: 0.9500, 65.13%, 4.26 toks/fwd,
against 0.9507 / 65.12% / 4.26 measured on the parent commit the same day.
Qwen3.5-27B tp2, the hybrid that exercises both counters: sizing identical at
`entry_bytes=266240`, and a bind probe on both arms agrees line for line on
all 64 layers -- shape, stride, dtype and storage offset -- so the mapping is
unchanged rather than merely close. GSM8K interleaved two runs per arm
overlaps (0.8767 / 0.8633 against 0.8333 / 0.8667); this configuration's
run-to-run spread is 3.3pp.

Not exercised: P/D, MiMo-V2, Kimi-K3.

ModelRunner carried four predicates over `model_type` -- `is_deepseek_mla`,
`is_qwen_next`, `is_kimi_linear`, `is_deepseek_v4` -- and set four booleans
from them, which the selector then re-composed into one answer through a chain
of ifs. The priority between them (V4 beats a latent rank, a hybrid beats one
too) existed only in the order of that chain, and the runner had to know
strings like `qwen3_5_moe_text` to feed it.

It is `Family`, one value, decided by `attn_family(hf_text_config)` next to the
dispatch it feeds. An enum rather than a string, so a name that does not exist
raises instead of falling through to MHA.

Two of the four predicates were also wrong, and the shape says so:

  - MLA is the presence of a latent rank, not membership of a list. The list
    named six model types and missed `k3_dspark`, whose config carries
    `kv_lora_rank=512`. That config reaches the engine only as a draft, and
    `draft_kv` already asked the right question -- `bool(kv_lora_rank)` -- so
    the mistake sat on a path nothing ran. It is one question now, asked the
    same way for a target and a draft, which also retires that fifth spelling.
  - An EAGLE draft is not an attention. `eagle` names a decoding algorithm and
    the old predicate branched on it, then hard-coded `deepseek_v2`/`deepseek_v3`
    for the stack underneath -- so an EAGLE trained on any other MLA base read
    as MHA. Two drafts of one target need not share a stack: the deployed
    MiniMax-M3 draft declares `model_type=llama`. Nothing about the algorithm
    appears here now.

`is_deepseek_v32` goes with them, renamed `has_mla_indexer`: every one of its
twelve readers asks whether this model's MLA rows carry an indexer key cache,
and GLM-5.2 is sparse too, so the name was already wrong for one of its own
users -- `aiter_mla` reassigned it to `self.is_sparse` before using it. One of
those readers was redundant and is gone: `forward_vars["sparse_kv_indptr"]`
exists exactly when the model has an indexer, so the buffer's own presence is
the test, which is how `eagle_proposer` had always spelled it.

`block_size` leaves the selector entirely. It was threaded through all three
functions and read by none.

Verified by dispatch parity over every config on the box -- family and indexer,
old spelling against new, 16 configs. Fifteen agree; the sixteenth is
`k3_dspark`, mha -> mla, on the path described above. Unit tests 5482 passed /
4 failed, the four being aiter-side failures already on HEAD, with 22 new ones
pinning the two corrections; reverting either makes them red.

GSM8K per family: MiniMax-M3 tp4 0.9356 at 73.24% acceptance (nine runs on this
branch span 0.9356-0.9462); DeepSeek-V4-Flash-DSpark tp2 0.9500, 64.86%, 4.24
toks/fwd against 0.9484/65.09%/4.25; GLM-5.2-MXFP4 tp4 0.9340, CI threshold
0.92 -- the only config on the box with an indexer, so the rename's riskiest
path; Qwen3.5-27B tp2 0.8533, inside the 0.8333-0.8767 its own noise spans.
Kimi-K3 could not be run: aiter-side breakage unrelated to this.

…once

Four hand-written byte formulas in the DeepSeek-V4 backend each described a
layout that was also described somewhere else, and nothing tied the copies
together:

  - how many planes this build has was answered independently five times --
    `rope_row_bytes` returning 0, two truthiness filters, `_kv_planes`'
    `is not None`, and an `if self._kv_fp8` in the carve. Only the last one
    fails loudly when it disagrees; the carve just builds one plane fewer.
  - what one indexer block costs was two formulas, one for sizing and one for
    allocation. They agreed by inspection, and a disagreement would have made
    sizing quietly wrong.
  - where the two regions sit inside that block was spelled out four times,
    the fourth being a scale view recomputing `rows * head_dim` from the pool
    shape. A view that starts one region early reads the tail of the data --
    valid-looking numbers, not a crash.
  - `head_dim` appeared in both the state window's shape and its row width,
    which have to agree and had nothing making them.

Each becomes one `EntryField` list in `pool_layout/v4_pool_fields.py`, the
sibling of `v4_pool_geometry`: that one says where rows go, this one says what
a row is. The backend keeps the policy -- the gfx fallback off fp8,
`--index_cache_dtype`, whether this GPU has an FP4 indexer -- and the shapes
follow from the answer with no branch of their own. A bf16 build declares one
plane instead of declaring two and setting one to zero bytes.

`indexer_block_regions` packs with nothing between regions rather than reusing
`entry_arena`'s aligned walk. Every consumer recomputes a region's start from
`rows` and the head dim rather than reading one, so borrowing the aligned walk
would have introduced exactly the gap that goes unnoticed. Both shipped
layouts land on 256 B anyway, which is why the tests scan head dims and pin
one case that does not.

Relaxes one existing invariant, deliberately: `pool_layout/` members may now
import a sibling member. The rule that matters is that a member is reachable
without aiter or the rest of atom so CI can import it, and a sibling is held
to that same rule by that same test -- nothing new is reached. Reading every
relative import as "leaves the tree" was the cheaper check and was already
charging for itself, in `entry_arena`'s note that a field list and an arena
over it could not be split apart. `test_a_sibling_import_names_a_member` is
what keeps the exemption to files the rule actually covers.

Verified: sub-pool bytes byte-for-byte unchanged on V4-Flash-DSpark tp2
dspark-5 (`kv 977152 / state 18652160`); GSM8K 1319 0.9522, DSpark acceptance
4.25-4.35 toks/fwd, which is the probe that fails first if the scale region
moved; unit tests 5527 passed with the four pre-existing failures; ruff
unchanged at 905.

…he walk

A pool row was a counter the bind walk advanced: `take_slot(kind)` on the
attention builders, `_next_layer_id` on the draft's own pool, and `layer_id`
itself where a dense model made the two agree. Three implementations of one
idea, and each of them assumed the runner's walk visits in the same order the
sizing walk did.

It is `pool_rows[kind][module]` now, assigned by one walk over
`_pooled_modules` that also answers how many rows to buy. Sizing and binding
read the same map rather than agreeing by convention, and the case that
convention did not cover is the reason: the runner's bind walk offers a draft's
modules to the target builder whenever a sibling builder declines them, and a
counter hands one the next row -- past the end of a pool sized without it --
without saying anything. A map raises.

With the row a property of the module rather than of the count, every
config-derived layer count that predicted the same thing goes:
`_get_total_num_layers`, `_num_draft_kv_layers`, `_num_cache_rows`,
`_owns_module`, `runner.num_full_attn`, `runner.num_gdn_attn_state`, and the
`num_kv_heads` / `num_draft_layers` parameters of `allocate_kv_cache_tensors`,
which no implementation read. Counting the modules is also what PP and a draft
stack needed those predictions to special-case.

`build_kv_cache_tensor`'s `layer_id` goes with them. It was the walk's ordinal
under a name that says otherwise, and two backends had already worked around
it -- DeepSeek-V4 re-deriving the layer id from the module prefix, MLA reading
`module.layer_num` and calling it `global_layer_id`. Two things it was standing
in for were wrong:

  - Kimi keyed `index_cache_layer_map` with it. That map is keyed by GLOBAL
    layer ids, which is why `aiter_mla` deliberately uses `module.layer_num`
    there; on a non-first PP stage local row 0 may be global layer 39.
  - the RapidServe import path keyed its whole `kv_cache_data` by
    `enumerate(...)` while the allocating path keyed by `module.layer_num`.
    One registry, two keyings, agreeing only while the first bound module was
    layer 0.

`PoolRowsMixin` lives in `pool_layout/`, which imports nothing: a draft's own
pool needs it too, and `spec_decode` promises in its own docstring to import no
attention.

Verified: unit tests 5530 passed with the four pre-existing failures; ruff
unchanged at 905. DeepSeek-R1-0528-MXFP4 MTP3 tp4 GSM8K 0.9439 (CI threshold
0.93) at 2.93-2.99 toks/fwd, binding 62 mla rows = 61 target + 1 draft, the
count the deleted arithmetic gave. Qwen3.5-27B sizing byte-identical to HEAD
(kv 877755 x 266240, state 256 x 39223296) and 48 + 16 rows, again the counts
the deleted formulas gave -- but its GSM8K read 0.85 against a same-session
HEAD control of 0.89, one run each, and that gap is not yet explained. K3 does
not start on this machine, so its half is covered by unit tests only.

`test_minimax_m3_index_topk` imports `index_topk`, which defines `@triton.jit`
kernels -- and a decorator runs at import. CI has no triton, so the import was
a collection error, and one of those takes the whole run down rather than the
one file. The tests now `importorskip("triton")`, which is a visible skip.

Their module also took triton from `vllm.triton_utils`, as does `sparse_attn`.
That is how they arrived from vLLM; the repository decided against it when the
same files were vendored into `fla_ops` -- "couples to vLLM internals. ATOM
uses plain triton imports throughout" -- and 73 files import it directly. These
two were the last that did not.

The skip means CI runs none of these, including the launch-policy tests that
need no device. Splitting the plain-Python half into a module of its own is
what would get them run there.

The PostToolUse hook ran `ruff check --fix` and captured its output without
ever reading it. Rules `--fix` cannot fix -- RUF012 is one it only advises on
-- therefore vanished. CI does not let them vanish: reviewdog filters to the
pull request's diff context, so a finding on a line an edit merely moved is a
red check whatever the repository's total was before, which is why counting
`ruff` findings before and after a change is not a gate for it.

Whatever survives `--fix` now goes back on stderr with exit 2, which for a
PostToolUse hook means "show this, the tool already ran". Pre-existing findings
in the edited file are reported too, on purpose: those are exactly the ones CI
picks up once an edit lands beside them.

Then everything it reported in the files this branch touches: SIM102, RUF012,
RUF013, RUF059, PLR1730, G201/TRY401, and a debug log for S112. Three BLE001
sites are `noqa` with a reason rather than narrowed -- each is a deliberate
best-effort probe of an optional dependency, and guessing which exception it
raises is how a probe becomes a startup crash.

Four assertions validate input rather than state: `max_block` follows from
`--max-model-len` and `--block-size`, and the two kpool checks read
`--block-size` against the model's `index_kpool`. Under `-O` they vanish, and
what each lets through is silent -- a block id wrapping into the packed
top-k's tie-break field, a truncating floor division, and a layout
`deepgemm_fp8_paged_mqa_logits` computes wrongly.

`_assert_packable` is `_require_packable` now, since it no longer asserts.

Also return 1 chunk when `_decode_score_chunks` has nothing to score. Capping
`chunks` is not enough: the round trip divides by its own inner `cdiv`, which
is zero there, and the raise lands outside any launch -- one rank down while
the others wait in the collective that follows.

Two ways a caller's buffer was accepted and then quietly misused.

`carve` sliced without checking the length. Slicing past the end truncates,
so a short buffer came back as a short last region and surfaced as an arena
complaining about bytes the caller never chose, in a place it did not slice.
It already computes the total; compare against it.

`LayerMajorArena` dropped a field's non-zero fill when handed a buffer. It
cannot apply one -- the buffer may be an imported pool already holding the
peer's KV -- but dropping it turns a `-inf` score plane into 0.0, which reads
as "always selected" instead of "never". Refuse instead. Inert today: no
`LayerMajorArena` field declares a fill.

`--block-size` and the block a backend's kernels index are two things tied
only by the first being a whole number of the second. They were one parameter
named `block_size` at every seam, so a draft resolving to another backend was
handed the target's answer instead of deriving its own.

The rule is now `AiterBackend.attn_block_size(hf_config, scheduler_block_size)`,
read by both the builder and `make_kv_pool`; the seam carries the scheduler's
block, which every model in the process shares, and each backend answers from
its own config. `MhaKvPool` shapes by exactly one token count -- its own block
-- plus `blocks_per_entry`, a count, because `entry_bytes` is what `page_pool`
charges. Charging per block instead would multiply the arena's 256 B field
alignment by the same factor: 9% of the pool at 8 blocks an entry, almost all
padding on the 64 B scale planes.

Names now say whether they are a size or a count: `entries` (not `blocks`) on
`pool_bytes`/`allocate`, `entry_tokens`, `blocks_per_entry`,
`scheduler_block_size`, `sparse_block_size`.

`TritonMHABackend` subclasses `AiterBackend` -- one pool, one layout, one
block rule, and the same `get_impl_cls`; the hand-written forwarding is gone.
Which backend a builder feeds is a `BACKEND` ClassVar rather than a named
class, so a subclass reading another backend's rule gets that one.

MiniMax-M3's indexer leaves `MhaKvPool`. It is `extra_fields`, declared by the
backend that knows what it holds and read back through `field_view` in the
shape only that backend knows; the pool is told bytes. `MhaKvPool` knows only
MHA again.

`adopt_imported_kv_pool` is deleted. It was `allocate_kv_cache_tensors` over
an imported buffer with one difference -- it did not publish what the hook
returns -- so the P/D decode side was left dereferencing runner attributes
nobody set. Both paths go through `ModelRunner._back_paged_pools` now.

`_local_draft_rows` counts a shared draft's MLA rows from the draft's own
modules. Deriving it as `pool rows - this stage's layer span` mixes units on a
hybrid (the span counts every layer, the pool only the MLA ones) and goes
negative, which empties the draft's half of the index-cache layout. Two
asserts state what a draft under PP is: built whole on the last stage, split
across none.

P/D regions carry a `semantic_role`. Folding the fields changed their order
-- `K0,V0,K1,V1...` became all of K then all of V -- at the same count and the
same bytes, so position alone lets a mismatched pair transfer K into V and say
nothing. The role is spelled out rather than interpolating `KvGeometry`, whose
repr would put class and field names on the wire.

The decode side of a P/D pair imports the paged pool as a handle and never
runs sizing, so `pool_plan` is empty there. Only the paged count arrives with
the handle; a STATE class has no entry count on that side at all, which means
`allocate_per_req_cache` cannot run and the attributes it publishes -- GDN's
`mamba_k_cache`, V4's arena, GLM-5.3's k-pool tail -- never exist.

This was always true. It used to surface as a hand-written bind loop that
silently skipped the linear-attention modules; since that loop was replaced by
the builders' own hook it surfaces as an `AttributeError` on whichever layer
binds first. Refuse by name instead, before anything is backed.

Read off the same `sub_pool_specs` declaration sizing reads, so a backend that
grows per-request state is covered without being listed, and the message names
the classes actually found rather than a table of families that would go
stale. `_has_state_pool` becomes `_state_pool_names` because both callers want
to say which.

`plan_field_planes` empties a plane whenever the fields fit in fewer, which
needs no unusual shape -- one field over two planes does it:

    widths [64, 64]  -> [['a'], []]
    widths [64, 512] -> [[], ['a']]

`DeepseekV4AttentionMetadataBuilder` hands each plane straight to
`EntryMajorArena`, which refused an empty field list, so a legal layout was a
startup crash waiting on a field count.

Allowed rather than guarded at the call site: a plane is addressed by index and
still costs its rows, so it has to exist even when it holds nothing -- dropping
it would shift every later plane. That is also why the same is not allowed of
`LayerMajorArena`, where an empty group is a region the model does not want and
`carve_layer_major` drops it to None.

The removed rejection test is replaced by four that state what such a plane is:
zero bytes an entry, no answer to `view`, and still spanning the stride its
rows were priced at.

…es it

`MhaKvPool.from_hf_config` took `layers=hf_config.num_hidden_layers` while
every row index came from `PoolRowsMixin`'s walk over the built model. Two
answers to one question, which is the split that mixin exists to close: one
layer too many is a layer of the pool paid for and never addressed, one too few
is a write past the end, and the startup byte check sees neither -- both sides
read the same inflated `entry_bytes`.

The count is the caller's now. A config says what one row holds; only the walk
knows how many rows there are.

That count is not knowable when `draft_kv_builder` runs -- it is called from
the proposer's `__init__`, which is `build_drafter` still running, so
`runner.drafter` is not set. So the builder is handed a way to build the pool
rather than the pool, and `kv_pool` builds it on first use; every reader is a
sizing or binding step and all of them run later. `invalidate_pool_rows` drops
the pool with the walk that sized it, so the two cannot part again.

Which left `make_kv_pool` answering two questions -- whether this flavor wants
a pool at all, and how big. The first is a property of the backend class and
never read `hf_config`, so it is `DRAFT_OWNS_KV_POOL` now and the factory is a
factory: the base raises instead of returning None. Verified the flag matches
what each of the eight backends used to answer.

No rows is a contradiction rather than a default: this builder exists because
the flavor owns a pool, and a pool of none prices at zero and binds nothing --
a draft running on no KV at all, with no error and no log. Indexing gets the
`KeyError` `pool_rows` promises.

Also moves `KvGeometry` to `pool_layout.pool_rows`, next to the walk that keys
by it. It is the one row space that is not a plain string, so the default
NamedTuple repr reached a startup log --

    caches 48 linear_state rows, 16 KvGeometry(num_kv_heads=2, head_dim=256) rows

-- and the `semantic_role` its regions carry. `__str__` gives both `h2d256`,
and the move puts it where a test can pin that without importing aiter.

Corrects what that role is for while here: the mooncake connector pairs the two
ends by list position and drops it, so it names a region in `set_block_count`'s
refusal and nothing more. The earlier claim that a peer matches by it was
wrong.

Each of these names something with zero occurrences left in `atom/`:

  - `test_state_arena.py` -> `test_entry_arena.py`, renamed with the module
  - `Eagle3DraftBuilder` -> `DraftKvBuilder`
  - `ModelRunner.is_qwen_next()` / `is_kimi_linear` -> `attn_family` resolving
    to `Family.GDN`, which is the one family predicate now. The point the
    comment was making survives: it reads the text config and the guard beside
    it reads the wrapper's `model_type`.
  - `runner.use_mla`, on a whole section of the architecture guide describing a
    branch `propose()` no longer has. Every step calls the builder's
    `prepare_mtp_decode` and merges what it returns, so which fields a step
    updates is the backend's answer; the one thing that does branch is
    `fuse_mtp_decode_position_update`, a builder capability rather than a model
    family.

And the FP4 MQA comment, which still explained a 512 floor that is 4096.

`token_to_seq_idxs` said which axis it indexed only by convention, and there
are two: query lengths give one id per new token, context lengths one per
gathered cached+new token. They are the same dtype and the same shape family,
so mixing them does not fault -- it resolves a token to another request's row
and attends on its behalf.

Both are spelled out now, after the `cu_seqlens_q` / `cu_seqlens_k` each
follows: `batch_id_per_q_token` and `batch_id_per_k_token`. The builder is
`build_batch_ids`, since it no longer belongs to one of them.

The k axis is new as a first-class field. It is layer-invariant, so the prefill
builder makes it once per forward instead of every layer's prefix gather
rebuilding it, and TBO slices it rather than rebuilding: the parent lays the
requests out contiguously and `cu_seqlens_k` says where a ubatch starts.

Which builder to use is decided by where the result must live, and that is
written down in the module rather than left to the call sites:
`build_batch_ids` covers any width bounded by `max_num_batched_tokens`, which a
pinned `CpuGpuBuffer` can hold, and every q-axis caller is one;
`build_batch_ids_device` covers the `total_kv` width, which nothing caps.

`build_batch_ids_device` takes the total as a required argument. Without it
`torch.repeat_interleave` synchronizes the device to read back its own output
length -- verified under `set_sync_debug_mode("error")` -- once per prefill,
between the metadata and the first layer. The caller already holds the sum as
`total_kv`; `context_lens` is `running_bs` wide but its padded tail is zero, so
the two agree. Required and not defaulted because omitting it is only slower,
which is the kind of thing a default hides.

`batch_id_per_token` survives in exactly two places, as aiter's own parameter
name, each marked so the next rename does not sweep it along.

Also clears ruff across every file touched here (188 findings, 152 of them
mechanical). Three that are not: `safetensors`' handle has no `__contains__`,
so SIM118's `not in h.keys()` rewrite would raise; the blind excepts are
probes whose fallback IS the answer, so they keep their behavior and carry a
reason; and the collapsible-`if` rewrites needed their bodies re-indented.

… the dict

`_build_v4_indexer_meta` returns two different dicts. The DECODE short-circuit
starts from `meta = {}` and adds only the FP4 window schedule; the prefill path
falls through to a literal with a dozen more fields. Renaming the row-to-
sequence map pointed `_score_topk_decode_fp4` at `batch_id_per_q_token`, which
only the prefill literal published -- so the first real FP4 decode raised
`KeyError`. Nothing in the unit suite reaches that forward; CI found it.

Republishing it on the decode branch would fix the symptom and keep the cause:
the dict entry was a copy of `attn_metadata.batch_id_per_q_token`, and
`indexer_meta` IS `attn_metadata.indexer_meta`, so both scorers already hold
the object it was copied from. Two branches copying one tensor is two places to
forget. Both read it off the metadata now and the entry is gone.

That also removes it from the two ubatch clone lists, where it was a second
clone of a tensor those paths already clone as an attribute -- the token-split
path at `ub_attn.batch_id_per_q_token`, the balanced path at the slice-and-
rebase that builds `ub.batch_id_per_q_token`.

The plugin bridges still write the key into their own indexer_meta dicts (six
sites across sglang and vllm). They are dead now, but those paths are neither
built nor tested here, so they are left for whoever owns them.

Gate: DeepSeek-V4-Flash-DSpark tp2 MTP3, GSM8K 1319 -> 0.9500 / 0.9507 against
a 0.9522 +/- 0.0059 baseline. gfx950 resolves `index_cache_dtype` to fp4 by
default, so the FP4 indexer decode path this fixes is the one that ran.

…he host

`prepare_mtp_decode` states its own contract in the first paragraph: eagle
leaves the CPU mirrors stale, so this routine rebuilds on the GPU and touches
no host staging. Routing the row -> sequence map through `publish_batch_ids`
broke that. It stages via the buffer's pinned host mirror, and
`_gate_staging_reuse` -- whose docstring names this very tensor as the one it
was written for -- fences that mirror once per FORWARD. A draft step runs
between two forwards, so its host write sits behind no fence and overwrites the
source of the verify forward's still-in-flight H2D.

The map goes back to a device copy out of the resident `row_ids` arange, which
is why that buffer had three readers and not two.

Only the MTP configurations regressed, which is the shape of the bug: the
non-speculative DeepSeek-V4-Pro, its DEP RCCL variant and DSpark all stayed
green, and DSpark's proposer never calls this routine. It also only shows under
CUDAGraph -- an eager run keeps the host close enough to the device that the
copy has landed -- which is why an `--enforce-eager` gate missed it.

DeepSeek-V4-Pro MTP3 tp8, CI's own arguments, GSM8K 1319, one variable:

  before  accept 50.95%  {0: 6.5, 1: 42.8, 2: 42.2, 3: 8.6}   0.9469 / 0.9454
  after   accept 66.00%  {0: 5.6, 1: 22.2, 2: 40.8, 3: 31.4}  0.9492 / 0.9500

against the last green CI run's 65.35% {6.7, 23.2, 37.7, 32.5} and 0.9530.
The acceptance distribution is the sharp end: a draft reading another request's
row lands its third token 9% of the time instead of 32%. CI's own
`mtp_accept_threshold` is 0.622.

Unit suite unchanged at 5587 passed; the 4 failures in
`test_deepseek_v4_wo_a_dequant.py` reproduce with this patch stashed.

…n rule

`make_kv_pool` took the block from `attn_block_size` on the draft's own
config, which is the wrong source: the pool it builds is read with the block
tables the target builder published, and the draft's kernels take the block off
the cache they were bound to. A block of its own addresses one page with an id
that counted another, and a draft has no metadata of its own to convert between
the two -- where the target's kernels index a finer block than the scheduler,
the target converts the ids itself, and there is no such place in a draft.

MiniMax-M3 + Eagle3 is where the two rules disagree, and there it could not
start at all: the target's block is the config's `sparse_block_size` (128)
whatever `--block-size` says, while the Eagle3 head is a plain Llama and gets
`scheduler if in (256, 1024) else 16`. `--block-size 128` failed the draft
check with 16 vs 128, no flag failed `model_runner.block_size % 128`, and 256
or 1024 moved both numbers without making them meet. No value existed.

The block is now the target builder's, read where the pool is built --
`DraftKvBuilder.kv_pool`, which is lazy, and that is what makes it reachable:
the builder is constructed after `build_drafter`.

`allocate_kv_cache_tensors`'s guard now compares against the SCHEDULER's block
instead of the target's, which the pool now equals by construction. That is not
the same check: `blocks` there counts scheduler blocks -- it is the budget
`paged_pool_bytes` charged per -- and one entry answers one of them, so the
block the pool took has to be that one too. It catches a target whose kernels
index a finer block than the scheduler, which is the case this could no longer
express: such a config used to start and mis-address, and now refuses.

Verified by starting MiniMax-M3-MXFP4 tp4 with the Eagle3 draft, which no
`--block-size` could do before. Its completions are NUL at the default
compilation level, but that is this model's own standing bug and not this
change: a single-variable `--level 0` control on the same server answers
`6*7` with `42`, and both arms block the draft's pool identically. Nothing
caught the breakage because the PR-level accuracy set has MiniMax-M3-MXFP4
without the draft, and the Eagle3 entry is nightly.

Unit suite 5588 passed, the 4 failures in `test_deepseek_v4_wo_a_dequant.py`
reproduce with this patch stashed. DeepSeek-V4 does not reach any of this:
`DeepseekV4Backend` inherits `DRAFT_OWNS_KV_POOL = False`, so `draft_kv_builder`
returns None and no pool is built here.